{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2016 Google LLC  \n",
    "  \n",
    " Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    " you may not use this file except in compliance with the License.  \n",
    " You may obtain a copy of the License at  \n",
    "  \n",
    "     http://www.apache.org/licenses/LICENSE-2.0  \n",
    "  \n",
    " Unless required by applicable law or agreed to in writing, software  \n",
    " distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    " WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    " See the License for the specific language governing permissions and  \n",
    " limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the dependency .py files, if any.\n",
    "! git clone https://github.com/GoogleCloudPlatform/cloudml-samples.git\n",
    "! cp cloudml-samples/census/estimator/trainer/* .\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Runs the training of the Wide & Deep model based on hyperparameter\n",
    "values received as input parameters.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import json\n",
    "import os\n",
    "import tensorflow as tf\n",
    "import input as input_module\n",
    "import model as model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_session_config_from_env_var():\n",
    "    \"\"\"Returns a tf.ConfigProto instance that has appropriate device_filters\n",
    "    set.\"\"\"\n",
    "\n",
    "    tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))\n",
    "\n",
    "    if (tf_config and 'task' in tf_config and 'type' in tf_config['task'] and\n",
    "            'index' in tf_config['task']):\n",
    "        # Master should only communicate with itself and ps\n",
    "        if tf_config['task']['type'] == 'master':\n",
    "            return tf.ConfigProto(device_filters=['/job:ps', '/job:master'])\n",
    "        # Worker should only communicate with itself and ps\n",
    "        elif tf_config['task']['type'] == 'worker':\n",
    "            return tf.ConfigProto(device_filters=[\n",
    "                '/job:ps',\n",
    "                '/job:worker/task:%d' % tf_config['task']['index']\n",
    "            ])\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_evaluate(args):\n",
    "    \"\"\"Run the training and evaluate using the high level API.\"\"\"\n",
    "\n",
    "    def train_input():\n",
    "        \"\"\"Input function returning batches from the training\n",
    "        data set from training.\n",
    "        \"\"\"\n",
    "        return input_module.input_fn(\n",
    "            args.train_files,\n",
    "            num_epochs=args.num_epochs,\n",
    "            batch_size=args.train_batch_size,\n",
    "            num_parallel_calls=args.num_parallel_calls,\n",
    "            prefetch_buffer_size=args.prefetch_buffer_size)\n",
    "\n",
    "    def eval_input():\n",
    "        \"\"\"Input function returning the entire validation data\n",
    "        set for evaluation. Shuffling is not required.\n",
    "        \"\"\"\n",
    "        return input_module.input_fn(\n",
    "            args.eval_files,\n",
    "            batch_size=args.eval_batch_size,\n",
    "            shuffle=False,\n",
    "            num_parallel_calls=args.num_parallel_calls,\n",
    "            prefetch_buffer_size=args.prefetch_buffer_size)\n",
    "\n",
    "    train_spec = tf.estimator.TrainSpec(\n",
    "        train_input, max_steps=args.train_steps)\n",
    "\n",
    "    exporter = tf.estimator.FinalExporter(\n",
    "        'census', input_module.SERVING_FUNCTIONS[args.export_format])\n",
    "    eval_spec = tf.estimator.EvalSpec(\n",
    "        eval_input,\n",
    "        steps=args.eval_steps,\n",
    "        exporters=[exporter],\n",
    "        name='census-eval')\n",
    "\n",
    "    run_config = tf.estimator.RunConfig(\n",
    "        session_config=_get_session_config_from_env_var())\n",
    "    run_config = run_config.replace(model_dir=args.job_dir)\n",
    "    print('Model dir %s' % run_config.model_dir)\n",
    "    estimator = model.build_estimator(\n",
    "        embedding_size=args.embedding_size,\n",
    "        # Construct layers sizes with exponential decay\n",
    "        hidden_units=[\n",
    "            max(2, int(args.first_layer_size * args.scale_factor**i))\n",
    "            for i in range(args.num_layers)\n",
    "        ],\n",
    "        config=run_config)\n",
    "\n",
    "    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARSER = argparse.ArgumentParser()\n",
    "# Input Arguments\n",
    "PARSER.add_argument(\n",
    "    '--train-files',\n",
    "    help='GCS file or local paths to training data',\n",
    "    nargs='+',\n",
    "    default='gs://cloud-samples-data/ml-engine/census/data/adult.data.csv')\n",
    "PARSER.add_argument(\n",
    "    '--eval-files',\n",
    "    help='GCS file or local paths to evaluation data',\n",
    "    nargs='+',\n",
    "    default='gs://cloud-samples-data/ml-engine/census/data/adult.test.csv')\n",
    "PARSER.add_argument(\n",
    "    '--job-dir',\n",
    "    help='GCS location to write checkpoints and export models',\n",
    "    default='/tmp/census-estimator')\n",
    "PARSER.add_argument(\n",
    "    '--num-parallel-calls',\n",
    "    help='Number of threads used to read in parallel the training and evaluation',\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--prefetch_buffer_size',\n",
    "    help='Naximum number of input elements that will be buffered when prefetching',\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--num-epochs',\n",
    "    help=\"\"\"\\\n",
    "  Maximum number of training data epochs on which to train.\n",
    "  If both --max-steps and --num-epochs are specified,\n",
    "  the training job will run for --max-steps or --num-epochs,\n",
    "  whichever occurs first. If unspecified will run for --max-steps.\\\n",
    "  \"\"\",\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--train-batch-size',\n",
    "    help='Batch size for training steps',\n",
    "    type=int,\n",
    "    default=40)\n",
    "PARSER.add_argument(\n",
    "    '--eval-batch-size',\n",
    "    help='Batch size for evaluation steps',\n",
    "    type=int,\n",
    "    default=40)\n",
    "PARSER.add_argument(\n",
    "    '--embedding-size',\n",
    "    help='Number of embedding dimensions for categorical columns',\n",
    "    default=8,\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--first-layer-size',\n",
    "    help='Number of nodes in the first layer of the DNN',\n",
    "    default=100,\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--num-layers',\n",
    "    help='Number of layers in the DNN',\n",
    "    default=4,\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--scale-factor',\n",
    "    help='How quickly should the size of the layers in the DNN decay',\n",
    "    default=0.7,\n",
    "    type=float)\n",
    "PARSER.add_argument(\n",
    "    '--train-steps',\n",
    "    help=\"\"\"\\\n",
    "  Steps to run the training job for. If --num-epochs is not specified,\n",
    "  this must be. Otherwise the training job will run indefinitely.\"\"\",\n",
    "    default=100,\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--eval-steps',\n",
    "    help='Number of steps to run evalution for at each checkpoint',\n",
    "    default=100,\n",
    "    type=int)\n",
    "PARSER.add_argument(\n",
    "    '--export-format',\n",
    "    help='The input format of the exported SavedModel binary',\n",
    "    choices=['JSON', 'CSV', 'EXAMPLE'],\n",
    "    default='JSON')\n",
    "PARSER.add_argument(\n",
    "    '--verbosity',\n",
    "    choices=['DEBUG', 'ERROR', 'FATAL', 'INFO', 'WARN'],\n",
    "    default='INFO')\n",
    "\n",
    "ARGUMENTS, _ = PARSER.parse_known_args()\n",
    "\n",
    "# Set python level verbosity\n",
    "tf.logging.set_verbosity(ARGUMENTS.verbosity)\n",
    "# Set C++ Graph Execution level verbosity\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = str(\n",
    "    tf.logging.__dict__[ARGUMENTS.verbosity] / 10)\n",
    "\n",
    "# Run the training job\n",
    "train_and_evaluate(ARGUMENTS)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
